#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LinPEAS Integration Module
LinPEAS集成模块
"""

import os
import re
import json
import time
import subprocess
from typing import Dict, List, Any, Optional
from pathlib import Path
from .tool_manager import ExternalToolManager

class LinPEASIntegration:
    """LinPEAS集成类"""
    
    def __init__(self, tool_manager: ExternalToolManager = None):
        self.tool_manager = tool_manager or ExternalToolManager()
        self.tool_name = "linpeas"
        self.last_scan_results = None
        self.scan_history = []
    
    def is_available(self) -> bool:
        """检查LinPEAS是否可用"""
        status = self.tool_manager.get_tool_status(self.tool_name)
        return status == "ready"
    
    def ensure_available(self) -> bool:
        """确保LinPEAS可用"""
        if self.is_available():
            return True
        
        print("📥 LinPEAS不可用，正在下载...")
        return self.tool_manager.download_tool(self.tool_name)
    
    def run_basic_scan(self, output_file: str = None) -> Dict[str, Any]:
        """运行基础扫描"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinPEAS不可用",
                "results": {}
            }
        
        args = ["-a"]  # 基础扫描参数
        
        if output_file:
            args.extend(["-o", output_file])
        
        return self._execute_scan(args, "basic")
    
    def run_fast_scan(self, output_file: str = None) -> Dict[str, Any]:
        """运行快速扫描"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinPEAS不可用",
                "results": {}
            }
        
        args = ["-f"]  # 快速扫描参数
        
        if output_file:
            args.extend(["-o", output_file])
        
        return self._execute_scan(args, "fast")
    
    def run_superfast_scan(self, output_file: str = None) -> Dict[str, Any]:
        """运行超快速扫描"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinPEAS不可用",
                "results": {}
            }
        
        args = ["-s"]  # 超快速扫描参数
        
        if output_file:
            args.extend(["-o", output_file])
        
        return self._execute_scan(args, "superfast")
    
    def run_network_scan(self, output_file: str = None) -> Dict[str, Any]:
        """运行网络扫描"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinPEAS不可用",
                "results": {}
            }
        
        args = ["-n"]  # 网络扫描参数
        
        if output_file:
            args.extend(["-o", output_file])
        
        return self._execute_scan(args, "network")
    
    def run_custom_scan(self, custom_args: List[str], scan_type: str = "custom", 
                       output_file: str = None) -> Dict[str, Any]:
        """运行自定义扫描"""
        if not self.ensure_available():
            return {
                "success": False,
                "error": "LinPEAS不可用",
                "results": {}
            }
        
        args = custom_args.copy()
        
        if output_file:
            args.extend(["-o", output_file])
        
        return self._execute_scan(args, scan_type)
    
    def _execute_scan(self, args: List[str], scan_type: str) -> Dict[str, Any]:
        """执行扫描"""
        print(f"🔍 开始LinPEAS {scan_type}扫描...")
        start_time = time.time()
        
        # 执行LinPEAS
        result = self.tool_manager.execute_tool(
            self.tool_name, 
            args, 
            timeout=600,  # 10分钟超时
            capture_output=True
        )
        
        execution_time = time.time() - start_time
        
        if result["success"]:
            # 解析结果
            parsed_results = self._parse_output(result["stdout"])
            
            scan_result = {
                "success": True,
                "scan_type": scan_type,
                "execution_time": execution_time,
                "raw_output": result["stdout"],
                "parsed_results": parsed_results,
                "timestamp": time.time(),
                "command": result["command"]
            }
            
            # 保存结果
            self.last_scan_results = scan_result
            self.scan_history.append(scan_result)
            
            print(f"✅ LinPEAS {scan_type}扫描完成 (耗时: {execution_time:.2f}s)")
            print(f"📊 发现 {len(parsed_results.get('findings', []))} 个发现")
            
            return scan_result
        else:
            error_result = {
                "success": False,
                "error": result["error"],
                "scan_type": scan_type,
                "execution_time": execution_time,
                "stderr": result["stderr"],
                "returncode": result["returncode"]
            }
            
            print(f"❌ LinPEAS {scan_type}扫描失败: {result['error']}")
            
            return error_result
    
    def _parse_output(self, output: str) -> Dict[str, Any]:
        """解析LinPEAS输出"""
        parsed = {
            "system_info": {},
            "users_info": {},
            "network_info": {},
            "processes_info": {},
            "files_info": {},
            "privilege_escalation": {},
            "findings": [],
            "recommendations": []
        }
        
        try:
            lines = output.split('\n')
            current_section = None
            
            for line in lines:
                line = line.strip()
                
                # 检测章节标题
                if self._is_section_header(line):
                    current_section = self._extract_section_name(line)
                    continue
                
                # 根据当前章节解析内容
                if current_section:
                    self._parse_section_content(line, current_section, parsed)
                
                # 检测重要发现
                finding = self._extract_finding(line)
                if finding:
                    parsed["findings"].append(finding)
                
                # 检测建议
                recommendation = self._extract_recommendation(line)
                if recommendation:
                    parsed["recommendations"].append(recommendation)
            
            # 统计信息
            parsed["statistics"] = {
                "total_findings": len(parsed["findings"]),
                "high_priority_findings": len([f for f in parsed["findings"] if f.get("priority") == "high"]),
                "medium_priority_findings": len([f for f in parsed["findings"] if f.get("priority") == "medium"]),
                "low_priority_findings": len([f for f in parsed["findings"] if f.get("priority") == "low"])
            }
            
        except Exception as e:
            print(f"⚠️ 输出解析警告: {e}")
            parsed["parse_error"] = str(e)
        
        return parsed
    
    def _is_section_header(self, line: str) -> bool:
        """检测是否为章节标题"""
        section_patterns = [
            r"═{10,}",  # 等号分隔符
            r"─{10,}",  # 横线分隔符
            r"\[.*\]",  # 方括号标题
            r"#{2,}",   # 井号标题
        ]
        
        for pattern in section_patterns:
            if re.search(pattern, line):
                return True
        
        return False
    
    def _extract_section_name(self, line: str) -> str:
        """提取章节名称"""
        # 移除特殊字符，提取核心内容
        clean_line = re.sub(r'[═─#\[\]]+', '', line).strip()
        
        # 常见章节映射
        section_mapping = {
            "system information": "system_info",
            "users": "users_info",
            "network": "network_info",
            "processes": "processes_info",
            "files": "files_info",
            "privilege escalation": "privilege_escalation"
        }
        
        for key, value in section_mapping.items():
            if key in clean_line.lower():
                return value
        
        return "general"
    
    def _parse_section_content(self, line: str, section: str, parsed: Dict[str, Any]):
        """解析章节内容"""
        if not line:
            return
        
        try:
            if section == "system_info":
                self._parse_system_info(line, parsed["system_info"])
            elif section == "users_info":
                self._parse_users_info(line, parsed["users_info"])
            elif section == "network_info":
                self._parse_network_info(line, parsed["network_info"])
            elif section == "processes_info":
                self._parse_processes_info(line, parsed["processes_info"])
            elif section == "files_info":
                self._parse_files_info(line, parsed["files_info"])
            elif section == "privilege_escalation":
                self._parse_privilege_escalation(line, parsed["privilege_escalation"])
        except Exception:
            pass  # 忽略解析错误
    
    def _parse_system_info(self, line: str, system_info: Dict[str, Any]):
        """解析系统信息"""
        if "hostname" in line.lower():
            system_info["hostname"] = line.split(":")[-1].strip()
        elif "kernel" in line.lower():
            system_info["kernel"] = line.split(":")[-1].strip()
        elif "distribution" in line.lower():
            system_info["distribution"] = line.split(":")[-1].strip()
    
    def _parse_users_info(self, line: str, users_info: Dict[str, Any]):
        """解析用户信息"""
        if "current user" in line.lower():
            users_info["current_user"] = line.split(":")[-1].strip()
        elif "sudo" in line.lower() and "users" in line.lower():
            if "sudo_users" not in users_info:
                users_info["sudo_users"] = []
            users_info["sudo_users"].append(line.strip())
    
    def _parse_network_info(self, line: str, network_info: Dict[str, Any]):
        """解析网络信息"""
        if "listening" in line.lower() and "port" in line.lower():
            if "listening_ports" not in network_info:
                network_info["listening_ports"] = []
            network_info["listening_ports"].append(line.strip())
    
    def _parse_processes_info(self, line: str, processes_info: Dict[str, Any]):
        """解析进程信息"""
        if "root" in line and "process" in line.lower():
            if "root_processes" not in processes_info:
                processes_info["root_processes"] = []
            processes_info["root_processes"].append(line.strip())
    
    def _parse_files_info(self, line: str, files_info: Dict[str, Any]):
        """解析文件信息"""
        if "writable" in line.lower():
            if "writable_files" not in files_info:
                files_info["writable_files"] = []
            files_info["writable_files"].append(line.strip())
    
    def _parse_privilege_escalation(self, line: str, priv_esc: Dict[str, Any]):
        """解析权限提升信息"""
        if "suid" in line.lower():
            if "suid_files" not in priv_esc:
                priv_esc["suid_files"] = []
            priv_esc["suid_files"].append(line.strip())
    
    def _extract_finding(self, line: str) -> Optional[Dict[str, Any]]:
        """提取发现"""
        # 检测高危发现
        high_risk_patterns = [
            r"password.*found",
            r"private.*key",
            r"suid.*root",
            r"sudo.*nopasswd"
        ]
        
        # 检测中危发现
        medium_risk_patterns = [
            r"writable.*config",
            r"world.*writable",
            r"interesting.*file"
        ]
        
        for pattern in high_risk_patterns:
            if re.search(pattern, line, re.IGNORECASE):
                return {
                    "type": "security_issue",
                    "priority": "high",
                    "description": line.strip(),
                    "category": "privilege_escalation"
                }
        
        for pattern in medium_risk_patterns:
            if re.search(pattern, line, re.IGNORECASE):
                return {
                    "type": "security_issue",
                    "priority": "medium",
                    "description": line.strip(),
                    "category": "file_permissions"
                }
        
        return None
    
    def _extract_recommendation(self, line: str) -> Optional[str]:
        """提取建议"""
        recommendation_patterns = [
            r"consider.*",
            r"recommend.*",
            r"should.*check",
            r"try.*"
        ]
        
        for pattern in recommendation_patterns:
            if re.search(pattern, line, re.IGNORECASE):
                return line.strip()
        
        return None
    
    def get_last_results(self) -> Optional[Dict[str, Any]]:
        """获取最后一次扫描结果"""
        return self.last_scan_results
    
    def get_scan_history(self) -> List[Dict[str, Any]]:
        """获取扫描历史"""
        return self.scan_history
    
    def export_results(self, file_path: str, format: str = "json"):
        """导出结果"""
        if not self.last_scan_results:
            print("❌ 没有可导出的结果")
            return False
        
        try:
            if format.lower() == "json":
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(self.last_scan_results, f, indent=2, ensure_ascii=False)
            
            elif format.lower() == "txt":
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write("LinPEAS扫描结果\n")
                    f.write("=" * 50 + "\n\n")
                    f.write(f"扫描类型: {self.last_scan_results['scan_type']}\n")
                    f.write(f"执行时间: {self.last_scan_results['execution_time']:.2f}秒\n")
                    f.write(f"时间戳: {time.ctime(self.last_scan_results['timestamp'])}\n\n")
                    f.write("原始输出:\n")
                    f.write("-" * 30 + "\n")
                    f.write(self.last_scan_results['raw_output'])
            
            print(f"📄 结果已导出到: {file_path}")
            return True
            
        except Exception as e:
            print(f"❌ 导出失败: {e}")
            return False
    
    def get_summary(self) -> Dict[str, Any]:
        """获取扫描摘要"""
        if not self.last_scan_results:
            return {"error": "没有可用的扫描结果"}
        
        parsed = self.last_scan_results.get("parsed_results", {})
        stats = parsed.get("statistics", {})
        
        return {
            "scan_type": self.last_scan_results["scan_type"],
            "execution_time": self.last_scan_results["execution_time"],
            "total_findings": stats.get("total_findings", 0),
            "high_priority": stats.get("high_priority_findings", 0),
            "medium_priority": stats.get("medium_priority_findings", 0),
            "low_priority": stats.get("low_priority_findings", 0),
            "recommendations_count": len(parsed.get("recommendations", [])),
            "timestamp": self.last_scan_results["timestamp"]
        }

# 使用示例
def example_usage():
    """使用示例"""
    # 创建LinPEAS集成实例
    linpeas = LinPEASIntegration()
    
    # 检查可用性
    if linpeas.is_available():
        print("✅ LinPEAS已就绪")
    else:
        print("📥 正在准备LinPEAS...")
        if not linpeas.ensure_available():
            print("❌ LinPEAS准备失败")
            return
    
    # 运行快速扫描
    print("🔍 开始快速扫描...")
    result = linpeas.run_fast_scan()
    
    if result["success"]:
        print("✅ 扫描完成")
        
        # 获取摘要
        summary = linpeas.get_summary()
        print(f"📊 扫描摘要:")
        print(f"  - 发现总数: {summary['total_findings']}")
        print(f"  - 高危发现: {summary['high_priority']}")
        print(f"  - 中危发现: {summary['medium_priority']}")
        print(f"  - 建议数量: {summary['recommendations_count']}")
        
        # 导出结果
        linpeas.export_results("linpeas_results.json", "json")
        
    else:
        print(f"❌ 扫描失败: {result['error']}")

if __name__ == "__main__":
    example_usage()